# Style 1: Only works with Kubernetes + common import options
from kubeflow.trainer import (
    Labels, Annotations, SpecLabels, SpecAnnotations,
    Name, PodTemplateOverrides, TrainerImage, TrainerCommand, TrainerArgs,
    PodTemplateOverride, PodTemplateSpecOverride, ContainerOverride
)

# Style 2: Import all options - SDK validates at runtime and rejects incompatible ones
from kubeflow.trainer.options import (
    # Kubernetes options
    Labels, Annotations, SpecLabels, SpecAnnotations, Name,
    PodTemplateOverrides, TrainerImage, TrainerCommand, TrainerArgs,
    # Common options (used in PodTemplateOverrides)
    PodTemplateOverride, PodTemplateSpecOverride, ContainerOverride,
    # LocalProcess options
    ProcessTimeout, WorkingDirectory
)

# Style 3: Backend-specific imports
from kubeflow.trainer.options.kubernetes import (
    Labels, Annotations, SpecLabels, SpecAnnotations, Name,
    PodTemplateOverrides, TrainerImage, TrainerCommand, TrainerArgs
)
from kubeflow.trainer.options.common import (
    PodTemplateOverride, PodTemplateSpecOverride, ContainerOverride
)
from kubeflow.trainer.options.localprocess import (
    ProcessTimeout, WorkingDirectory
)

# TrainJob resource metadata (for TrainJob level Labels and annotations)
Labels(labels={"team": "ml-platform", "project": "llm"})
Annotations(annotations={"description": "Fine-tuning job", "owner": "user"})
Name(name="my-custom-job-name")

# Derivative JobSet and Jobs labels (for Pod level Labels and annotations)
SpecLabels(labels={"app": "training", "version": "v1.0", "env": "prod"})
SpecAnnotations(annotations={"prometheus.io/scrape": "true", "cost-center": "ai"})

TrainerImage(image="custom/pytorch:2.0-cuda11.8")
TrainerCommand(command=["python", "train.py"])
TrainerArgs(args=["--epochs", "100", "--lr", "0.001", "--batch-size", "32"])

PodTemplateOverrides(pod_template_overrides=[
    PodTemplateOverride(
        target_jobs=["node", "launcher"],
        metadata={  # Optional: pod-level labels/annotations
            "labels": {"tier": "training"},
            "annotations": {"sidecar.istio.io/inject": "false"}
        },
        spec=PodTemplateSpecOverride(
            service_account_name="training-sa",
            node_selector={"gpu": "true", "zone": "us-west1"},
            affinity={
                "nodeAffinity": {
                    "requiredDuringSchedulingIgnoredDuringExecution": {
                        "nodeSelectorTerms": [{
                            "matchExpressions": [{
                                "key": "gpu-type",
                                "operator": "In",
                                "values": ["nvidia-a100"]
                            }]
                        }]
                    }
                }
            },
            tolerations=[
                {"key": "gpu", "operator": "Exists", "effect": "NoSchedule"},
                {"key": "high-priority", "operator": "Equal", "value": "true", "effect": "NoSchedule"}
            ],
            volumes=[
                {"name": "data", "emptyDir": {}},
                {"name": "model", "persistentVolumeClaim": {"claimName": "model-pvc"}}
            ],
            containers=[
                ContainerOverride(
                    name="node",
                    env=[
                        {"name": "NCCL_DEBUG", "value": "INFO"},
                        {"name": "WANDB_API_KEY", "valueFrom": {"secretKeyRef": {
                            "name": "wandb-secret", "key": "api-key"
                        }}}
                    ],
                    volume_mounts=[
                        {"name": "data", "mountPath": "/data"},
                        {"name": "model", "mountPath": "/model"}
                    ]
                )
            ],
            init_containers=[
                ContainerOverride(
                    name="init-setup",
                    env=[{"name": "SETUP_MODE", "value": "fast"}],
                    volume_mounts=[{"name": "data", "mountPath": "/setup"}]
                )
            ],
            scheduling_gates=[{"name": "wait-for-quota"}],
            image_pull_secrets=[{"name": "registry-credentials"}]
        )
    )
])

from kubeflow.trainer.options import ProcessTimeout, WorkingDirectory
ProcessTimeout(timeout_seconds=3600)
WorkingDirectory(working_dir="/custom/training/path")

job_name = client.train(
    runtime=client.get_runtime("torch-distributed"),
    trainer=CustomTrainer(
        func=lambda: print("Training LLM!"),
        num_nodes=2,
        resources_per_node={
            "cpu": "8",
            "memory": "32Gi",
            "nvidia.com/gpu": "2"
        },
        packages_to_install=["torch==2.0.0", "transformers", "datasets"],
        env={"WANDB_PROJECT": "llm-training"}
    ),
    options=[
        # TrainJob resource metadata
        Name(name="llm-finetune-001"),
        Labels(labels={"owner": "ml-team", "project": "llm"}),
        Annotations(annotations={"description": "GPT fine-tuning experiment"}),
        
        # Derivative JobSet/Jobs labels (for training pods)
        SpecLabels(labels={"app": "training", "version": "v1.0", "env": "prod"}),
        SpecAnnotations(annotations={"prometheus.io/scrape": "true"}),
        
        # Trainer container customization
        TrainerImage(image="my-registry/pytorch:2.0-cuda11.8"),
        
        # Pod-level customizations
        PodTemplateOverrides(pod_template_overrides=[
            PodTemplateOverride(
                target_jobs=["node"],
                spec=PodTemplateSpecOverride(
                    service_account_name="gpu-training-sa",
                    node_selector={"accelerator": "nvidia-a100", "zone": "us-west1"},
                    tolerations=[
                        {"key": "nvidia.com/gpu", "operator": "Exists", "effect": "NoSchedule"}
                    ],
                    volumes=[
                        {"name": "data", "persistentVolumeClaim": {"claimName": "training-data-pvc"}},
                        {"name": "cache", "emptyDir": {"sizeLimit": "50Gi"}}
                    ],
                    containers=[
                        ContainerOverride(
                            name="node",
                            env=[
                                {"name": "NCCL_DEBUG", "value": "INFO"},
                                {"name": "NCCL_IB_DISABLE", "value": "1"},
                                {"name": "WANDB_API_KEY", "valueFrom": {
                                    "secretKeyRef": {"name": "wandb-creds", "key": "api-key"}
                                }}
                            ],
                            volume_mounts=[
                                {"name": "data", "mountPath": "/data"},
                                {"name": "cache", "mountPath": "/cache"}
                            ]
                        )
                    ]
                )
            )
        ])
    ]
)